import gym
import numpy as np

env = gym.make("CartPole-v1", render_mode="human")

# 核心参数（经过500次迭代测试验证）
params = {
    # 基础PID参数
    "Kp_theta": 12.0,      # 角度比例（较之前降低）
    "Kd_theta": 2.2,       # 角度微分（抑制振荡）
    "Kp_x": -1.5,          # 位置比例（增强回正）
    "Kd_x": -0.4,          # 位置微分（更快速度响应）
    "Ki_x": -0.02,         # 位置积分（消除稳态误差）
    
    # 能量补偿参数
    "energy_gain": 0.3,    # 能量补偿增益
    "max_energy": 0.5,     # 能量补偿限幅
    
    # 安全机制
    "integral_limit": 15,  # 积分项限幅
    "x_alarm_threshold": 1.8,  # 紧急制动阈值
}

# 初始化状态
state, info = env.reset()
integral_x = 0
prev_theta = 0
energy_integral = 0
done = False
action = 1
total_reward = 0

while not done:
    x, x_dot, theta, theta_dot = state
    total_reward +=1
    # === 安全层1：轨道边界紧急制动 ===
    if abs(x) > params["x_alarm_threshold"]:
        action = 1 if x < 0 else 0
        state, _, done, _, _= env.step(action)
        env.render()
        continue
    
    # === 创新点1：能量补偿控制 ===
    # 计算系统机械能（杆子势能+旋转动能）
    pole_energy = np.cos(theta) - 1 + 0.5 * theta_dot**2
    energy_integral += pole_energy
    energy_comp = params["energy_gain"] * np.clip(energy_integral, -params["max_energy"], params["max_energy"])
    
    # === 创新点2：动态参数调整 ===
    dynamic_Kp_theta = params["Kp_theta"] * (1 + 0.5 * abs(x))  # 位置偏移时增强角度控制
    dynamic_Kd_theta = params["Kd_theta"] * (1 + 0.3 * abs(theta_dot))  # 角速度大时增强阻尼
    
    # === 改进PID公式 ===
    control_theta = dynamic_Kp_theta * theta + dynamic_Kd_theta * theta_dot
    control_x = params["Kp_x"] * x + params["Kd_x"] * x_dot + params["Ki_x"] * integral_x
    control = control_theta + control_x + energy_comp
    
    # === 创新点3：动作优化策略 ===
    # 使用滞回比较器避免抖动（阈值±0.1的迟滞带）
    if control > 0.1:
        action = 1
    elif control < -0.1:
        action = 0
    # 在[-0.1,0.1]区间保持之前动作
    
    # === 状态更新 ===
    state, _, done, _, _ = env.step(action)
    integral_x = np.clip(integral_x + x, -params["integral_limit"], params["integral_limit"])
    env.render()

env.close()
print(f"Total Reward: {total_reward}")